#include"memory.h"
#include"print.h"
#include"global.h"
#include"string.h"
#include"thread.h"
#include"debug.h"
#include"list.h"

#define KERNAL_USED 1024*1024
#define PDE_INDEX(vaddr) ((vaddr&0xffc00000)>>22)
#define PTE_INDEX(vaddr) ((vaddr&0x3ff000)>>12)

physical_addr_pool kernal_paddr_pool,user_paddr_pool;

virtual_addr_pool kernal_vaddr_pool;

mem_bank kernal_mm_bank;

#define BLOCK_DES_COUNT 7
void mem_bank_init(mem_bank* bank){
	uint32_t i;
	for(i=0;i<BLOCK_DES_COUNT;i++){
		bank->mem_block_des_vec[i].block_size=16<<i;
		bank->mem_block_des_vec[i].block_count_per_arena= \
	  	(PG_SIZE-sizeof(mm_meta))/bank->mem_block_des_vec[i].block_size;
		list_init(&bank->mem_block_des_vec[i].free_block_list);
	}
}

static mm_meta* block2meta(mem_block* block){
	return (mm_meta*)((uint32_t)block & 0xfffff000);
}

static mem_block* block_at_index(mm_meta* pm,uint32_t index){
	return (mem_block*)((uint32_t)pm+sizeof(mm_meta)+index*pm->block_des->block_size);
}

static inline virtual_addr_pool* vpool(p_type ptype){
	task_tcb* c=current_thread();
	return ptype==PG_K ? &kernal_vaddr_pool : &c->process_vaddr_pool;
}

static inline physical_addr_pool* ppool(p_type ptype){
	return ptype==PG_K ? &kernal_paddr_pool : &user_paddr_pool;
}

static inline p_type curr_thread_type(void){
	task_tcb* c=current_thread();
	return c->page_dir==NULL ? PG_K : PG_U;
}

static inline mem_bank* mm_bank(p_type ptype){
	task_tcb* c=current_thread();
	return ptype==PG_K ? &kernal_mm_bank : &c->process_mm_bank;
}

static mem_block_des* find_mm_des(uint32_t size){
	int i;
	mem_bank* mmb=mm_bank(curr_thread_type());
	for(i=0;i<BLOCK_DES_COUNT;i++){
		if(size<mmb->mem_block_des_vec[i].block_size)
			break;
	}
	return &mmb->mem_block_des_vec[i==BLOCK_DES_COUNT ? BLOCK_DES_COUNT-1 : i];
}

static void* make_large_mm_block(uint32_t size){
	uint32_t page_cnt=div_round_up(size,PG_SIZE);
	void* pp=get_page(curr_thread_type(),page_cnt);
	mm_meta* pm=pp;
	pm->large=true;
	pm->free_block_count=page_cnt;
	return (void*)(pm+sizeof(mm_meta));
}

static void make_sm_mm_block(mem_block_des* mmdes){
	uint32_t i;
	mem_block* pb;
	void* pp=get_page(curr_thread_type(),1);
	mm_meta* pm=pp;
	pm->block_des=mmdes;
	pm->large=false;
	pm->free_block_count=mmdes->block_count_per_arena;
	for(i=0;i<mmdes->block_count_per_arena;i++){
		pb=block_at_index(pm,i);
		list_append(&mmdes->free_block_list,&pb->b_tag);
	}
}

static void* find_mm_block(mem_block_des* mmdes){
	ASSERT(!list_empty(&mmdes->free_block_list));
	return list_pop(&mmdes->free_block_list);
}

void* sys_malloc(uint32_t size){
	bool large=size>1024 ? true : false;
	mem_block_des* mmdes=find_mm_des(size);
	if(large){
		return make_large_mm_block(size);
	}else{
		if(list_empty(&mmdes->free_block_list)){
			make_sm_mm_block(mmdes);
		}
	}
	return find_mm_block(mmdes);
}

void sys_free(void* addr){
	ASSERT(addr);
 	mm_meta* pm=block2meta(addr);
	p_type ptype=curr_thread_type();
	if(pm->large){
		free_page(ptype,(void*)pm,pm->free_block_count);
	}else{
		if(pm->free_block_count+1==pm->block_des->block_count_per_arena){
			free_page(ptype,(void*)pm,1);
		}else{
			list_append(&pm->block_des->free_block_list,(list_elem*)addr);
			pm->free_block_count++;
		}
	}
}

static void mem_physical_pool_init(uint32_t mem_size){
	//memory has allocate for pde and pte
	uint32_t free;	
	free=mem_size-KERNAL_USED-PG_SIZE*256;
	kernal_paddr_pool.paddr_bitmap.bits=(uint8_t*)PHY_BITMAP_BASE;
	kernal_paddr_pool.paddr_bitmap.bmap_bytes_len=free/PG_SIZE/4;
	kernal_paddr_pool.paddr_start=KERNAL_USED+PG_SIZE*256;
	bitmap_init(&kernal_paddr_pool.paddr_bitmap);

	//physical memory kernal and user half to half
	user_paddr_pool.paddr_bitmap.bits=(uint8_t*)(PHY_BITMAP_BASE+kernal_paddr_pool.paddr_bitmap.bmap_bytes_len);
	user_paddr_pool.paddr_bitmap.bmap_bytes_len=free/PG_SIZE/4;
	user_paddr_pool.paddr_start=kernal_paddr_pool.paddr_start+free/2;
	bitmap_init(&user_paddr_pool.paddr_bitmap);	
	
	//kernal virtual address pool init	
	kernal_vaddr_pool.vaddr_start=KERNAL_HEAP_START;
	kernal_vaddr_pool.vaddr_bitmap.bits=(uint8_t*)((uint32_t)user_paddr_pool.paddr_bitmap.bits+user_paddr_pool.paddr_bitmap.bmap_bytes_len);
	kernal_vaddr_pool.vaddr_bitmap.bmap_bytes_len=kernal_paddr_pool.paddr_bitmap.bmap_bytes_len;
	bitmap_init(&kernal_vaddr_pool.vaddr_bitmap);
	put_str("ALL MEM:");
	put_int(mem_size);
	put_char('\n');	
	put_str("FREE MEM:");
	put_int(free);
	put_char('\n');
	put_str("user PHY:");
	put_int(user_paddr_pool.paddr_start);
	put_char('\n');
	put_str("KHEAP:");
	put_int(kernal_vaddr_pool.vaddr_start);
	put_char('\n');
}

//return virtual address point to vaddr pde
void* pde_ptr(uint32_t vaddr){
	return (void*)(0xfffff000 | PDE_INDEX(vaddr)*4);
}

//return virtual address point to vaddr pte
void* pte_ptr(uint32_t vaddr){
	return (void*)(0xffc00000 | PDE_INDEX(vaddr)<<12 | PTE_INDEX(vaddr)*4);
}

void* vaddr2paddr(uint32_t vaddr){
	void* pte_entry=pte_ptr(vaddr);
	return (void*)(((uint32_t)(*(uint32_t*)pte_entry) & 0xfffff000) + (vaddr & 0xfff));
}

//malloc physical page for caller
static void* pmalloc(p_type ptype,uint32_t cnt){
	uint32_t i=0;
	physical_addr_pool* phy_pool=ppool(ptype);
	int32_t page_start_index=bitmap_scan(&phy_pool->paddr_bitmap,cnt);
	if(page_start_index<0)
		return NULL;
	while(i<cnt){
		bitmap_set(&phy_pool->paddr_bitmap,page_start_index+i,1);
		i++;
	}
	return (void*)(phy_pool->paddr_start+page_start_index*PG_SIZE);
}

//free physical page
void pfree(p_type ptype,uint32_t paddr,uint32_t cnt){
	uint32_t i=0;
	uint32_t page_index=0;
	physical_addr_pool* phy_pool=ppool(ptype);
	page_index=(paddr-phy_pool->paddr_start)/PG_SIZE;
	while(i<cnt){
		bitmap_set(&phy_pool->paddr_bitmap,page_index+i,0);
		i++;
	}
}

static void* mem_map(p_type ptype,uint32_t vaddr,uint32_t paddr){
	void* pde=pde_ptr(vaddr);
	void* pte;
	void* p;
	if(!*((uint32_t*)pde)){
		//page directory entry is not exist alloc a page for page table
		p=pmalloc(ptype,1);
		if(!p)
			return NULL;		
		*((uint32_t*)pde)=(uint32_t)p | PG_US_U | PG_RW_W | PG_P_1;
	}
	pte=pte_ptr(vaddr);
	/*
	 * here we add the page physical address into 
	 * page table entry which the virtual address 
	 * will be map to
	 */
	*((uint32_t*)pte)=paddr | PG_US_U | PG_RW_W | PG_P_1;
	return (void*)vaddr;
}

static void mem_unmap(uint32_t addr){
	void* pte=pte_ptr(addr);
	ASSERT(*(uint32_t*)pte & PG_P_1);
	*(uint32_t*)pte &=~PG_P_1;
	asm volatile("invlpg %0"::"m"(addr));
}

void* get_page(p_type ptype,uint32_t cnt){
	task_tcb* curr=current_thread();
	//user process must has alone page dir
	if(ptype==PG_U && curr->page_dir==NULL){
		PANIC("can't allocate memory \n");
	}
	//first get cnt page virtual page
	virtual_addr_pool* vp=vpool(ptype);
	void* vaddr;
	int32_t v_page_index=bitmap_scan(&vp->vaddr_bitmap,cnt);
	uint32_t i=0;
	if(v_page_index<0)
		return NULL;
	while(i<cnt){
		//second get one physical page
		void* paddr=pmalloc(ptype,1);
		if(!paddr){
			return NULL;
		}
		bitmap_set(&vp->vaddr_bitmap,v_page_index+i,1);
		//map virtual address to physicall address
		vaddr=mem_map(ptype,vp->vaddr_start+(v_page_index+i)*PG_SIZE,(uint32_t)paddr);
		if(!vaddr){
			return NULL;
		}
		memset(vaddr,0,PG_SIZE);
		i++;
	}
	return (void*)(vp->vaddr_start+v_page_index*PG_SIZE);
}

void free_page(p_type ptype,void* addr,uint32_t cnt){
	ASSERT(((uint32_t)addr & 0xfff)==0);
	uint32_t _addr=(uint32_t)addr;
	uint32_t i=0;
	uint32_t page_index=0;
	virtual_addr_pool* vp=vpool(ptype);
	page_index=(_addr-vp->vaddr_start)/PG_SIZE;
	while(i<cnt){
		bitmap_set(&vp->vaddr_bitmap,page_index+i,0);
		pfree(ptype,_addr+i*PG_SIZE,1);
		mem_unmap(_addr+i*PG_SIZE);
		i++;
	}
}

//map vaddr page to physical page
void* get_fix_vpage(p_type ptype,uint32_t vaddr){
	virtual_addr_pool* vp=vpool(ptype);
	uint32_t bits_index;
	bits_index=(vaddr-vp->vaddr_start)/PG_SIZE;
	void* paddr;
	paddr=pmalloc(ptype,1);
	if(!paddr){
		return NULL;
	}
	bitmap_set(&vp->vaddr_bitmap,bits_index,1);
	//page dirctory table always allocate in kernal memory
	mem_map(PG_K,vp->vaddr_start+bits_index*PG_SIZE,(uint32_t)paddr);
	memset((void*)vaddr,0,PG_SIZE);
	return (void*)vaddr;
}

void* get_fix_clone_vpage(p_type ptype,uint32_t vaddr){
	void* paddr;
	paddr=pmalloc(ptype,1);
	if(!paddr){
		return NULL;
	}
	mem_map(ptype,vaddr,(uint32_t)paddr);
	memset((void*)vaddr,0,PG_SIZE);
	return (void*)vaddr;
}

void mem_init(void){
	put_str("memory init\n");
	mem_bank_init(&kernal_mm_bank);
	mem_physical_pool_init(MEM_SIZE);
}
